"""
技术指标计算模块

使用pandas和numpy计算各种技术指标，并将结果保存到数据库
"""
import logging
import pandas as pd
import numpy as np
from sqlalchemy.orm import Session
from sqlalchemy import func

from modules.models.base import get_db
from modules.models.stock import SecurityInfo, DailyQuote, TechnicalIndicator

# 配置日志
logger = logging.getLogger("indicators_calculator")


class IndicatorsCalculator:
    """技术指标计算类"""
    
    def __init__(self):
        """初始化计算器"""
        self.db = next(get_db())
    
    def __del__(self):
        """析构函数，确保关闭数据库连接"""
        if hasattr(self, 'db') and self.db:
            self.db.close()
    
    def calculate_security_indicators(self, security_code, recalculate=False):
        """
        计算指定证券的技术指标
        
        Args:
            security_code: 证券代码
            recalculate: 是否重新计算所有指标，默认为False（仅计算新数据）
        
        Returns:
            计算的指标数量
        """
        # 获取证券信息
        security = self.db.query(SecurityInfo).filter(
            SecurityInfo.code == security_code
        ).first()
        
        if not security:
            logger.error(f"证券代码 {security_code} 不存在")
            return 0
        
        # 获取日线数据，按日期排序
        query = self.db.query(DailyQuote).filter(
            DailyQuote.security_id == security.id
        ).order_by(DailyQuote.trade_date)
        
        # 如果不是重新计算，只处理没有指标的数据
        if not recalculate:
            # 找出已经有技术指标的日期
            existing_quote_ids = {
                quote_id[0] for quote_id in 
                self.db.query(TechnicalIndicator.daily_quote_id).all()
            }
            
            # 过滤出没有技术指标的数据
            query = query.filter(~DailyQuote.id.in_(existing_quote_ids))
        
        # 检查是否有需要计算的数据
        count = query.count()
        if count == 0:
            logger.info(f"证券 {security_code} 没有需要计算指标的新数据")
            return 0
        
        # 获取所有日线数据并转换为DataFrame
        quotes = query.all()
        
        # 日线数据过少无法计算指标
        if len(quotes) < 30:
            logger.warning(f"证券 {security_code} 数据量不足，需要至少30条数据才能计算指标")
            return 0
        
        # 构建DataFrame
        data = {
            'id': [q.id for q in quotes],
            'date': [q.trade_date for q in quotes],
            'open': [q.open for q in quotes],
            'high': [q.high for q in quotes],
            'low': [q.low for q in quotes],
            'close': [q.close for q in quotes],
            'volume': [q.volume for q in quotes]
        }
        df = pd.DataFrame(data)
        
        # 计算技术指标
        indicators_df = self._calculate_indicators(df)
        
        # 保存到数据库
        saved_count = self._save_indicators(indicators_df)
        
        logger.info(f"成功计算并保存 {security_code} 的 {saved_count} 条技术指标数据")
        return saved_count
    
    def _calculate_indicators(self, df):
        """计算各种技术指标"""
        try:
            # 确保价格数据是数值类型
            price_columns = ['open', 'high', 'low', 'close']
            for col in price_columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
            
            df['volume'] = pd.to_numeric(df['volume'], errors='coerce')
            
            # 计算移动平均线
            df['ma5'] = df['close'].rolling(window=5).mean()
            df['ma10'] = df['close'].rolling(window=10).mean()
            df['ma20'] = df['close'].rolling(window=20).mean()
            df['ma30'] = df['close'].rolling(window=30).mean()
            df['ma60'] = df['close'].rolling(window=60).mean()
            df['ma120'] = df['close'].rolling(window=120).mean()
            df['ma250'] = df['close'].rolling(window=250).mean()
            
            # 计算成交量移动平均线
            df['vol5'] = df['volume'].rolling(window=5).mean()
            df['vol10'] = df['volume'].rolling(window=10).mean()
            
            # 计算RSI指标
            def calculate_rsi(prices, period=14):
                deltas = np.diff(prices)
                seed = deltas[:period+1]
                up = seed[seed >= 0].sum()/period
                down = -seed[seed < 0].sum()/period
                rs = up/down if down != 0 else 0
                rsi = np.zeros_like(prices)
                rsi[:period] = 100. - 100./(1. + rs)
                
                for i in range(period, len(prices)):
                    delta = deltas[i-1]
                    if delta > 0:
                        upval = delta
                        downval = 0.
                    else:
                        upval = 0.
                        downval = -delta
                    
                    up = (up * (period - 1) + upval) / period
                    down = (down * (period - 1) + downval) / period
                    
                    rs = up/down if down != 0 else 0
                    rsi[i] = 100. - 100./(1. + rs)
                
                return rsi
            
            df['rsi6'] = calculate_rsi(df['close'].values, period=6)
            df['rsi12'] = calculate_rsi(df['close'].values, period=12)
            df['rsi24'] = calculate_rsi(df['close'].values, period=24)
            
            # 计算MACD指标
            def calculate_macd(prices, fast=12, slow=26, signal=9):
                ema_fast = pd.Series(prices).ewm(span=fast, adjust=False).mean()
                ema_slow = pd.Series(prices).ewm(span=slow, adjust=False).mean()
                macd = ema_fast - ema_slow
                macd_signal = macd.ewm(span=signal, adjust=False).mean()
                macd_hist = macd - macd_signal
                return macd, macd_signal, macd_hist
            
            df['macd'], df['macd_signal'], df['macd_hist'] = calculate_macd(df['close'].values)
            
            # 计算布林带
            def calculate_bollinger_bands(prices, window=20, num_std=2):
                rolling_mean = pd.Series(prices).rolling(window=window).mean()
                rolling_std = pd.Series(prices).rolling(window=window).std()
                upper_band = rolling_mean + (rolling_std * num_std)
                lower_band = rolling_mean - (rolling_std * num_std)
                return upper_band, rolling_mean, lower_band
            
            df['boll_upper'], df['boll_middle'], df['boll_lower'] = calculate_bollinger_bands(df['close'].values)
            
            # 计算KDJ指标
            def calculate_kdj(high, low, close, n=9, m1=3, m2=3):
                high_n = pd.Series(high).rolling(window=n).max()
                low_n = pd.Series(low).rolling(window=n).min()
                rsv = 100 * ((pd.Series(close) - low_n) / (high_n - low_n))
                
                k = np.zeros_like(close)
                d = np.zeros_like(close)
                j = np.zeros_like(close)
                
                for i in range(n, len(close)):
                    if i == n:
                        k[i] = 50
                        d[i] = 50
                    else:
                        k[i] = (m1 * k[i-1] + (100 - m1) * rsv[i]) / 100
                        d[i] = (m2 * d[i-1] + (100 - m2) * k[i]) / 100
                    
                    j[i] = 3 * k[i] - 2 * d[i]
                
                return k, d, j
            
            df['kd_k'], df['kd_d'], df['kd_j'] = calculate_kdj(df['high'].values, df['low'].values, df['close'].values)
            
            # 构建结果DataFrame
            result = pd.DataFrame({
                'daily_quote_id': df['id'],
                'ma5': df['ma5'],
                'ma10': df['ma10'],
                'ma20': df['ma20'],
                'ma30': df['ma30'],
                'ma60': df['ma60'],
                'ma120': df['ma120'],
                'ma250': df['ma250'],
                'vol5': df['vol5'],
                'vol10': df['vol10'],
                'rsi6': df['rsi6'],
                'rsi12': df['rsi12'],
                'rsi24': df['rsi24'],
                'macd': df['macd'],
                'macd_signal': df['macd_signal'],
                'macd_hist': df['macd_hist'],
                'boll_upper': df['boll_upper'],
                'boll_middle': df['boll_middle'],
                'boll_lower': df['boll_lower'],
                'kd_k': df['kd_k'],
                'kd_d': df['kd_d'],
                'kd_j': df['kd_j']
            })
            
            return result
            
        except Exception as e:
            logger.error(f"计算技术指标失败: {str(e)}")
            raise
    
    def _save_indicators(self, df):
        """保存技术指标到数据库"""
        if df.empty:
            return 0
        
        try:
            # 批量处理
            batch_size = 100
            total_rows = len(df)
            saved_count = 0
            
            for start_idx in range(0, total_rows, batch_size):
                end_idx = min(start_idx + batch_size, total_rows)
                batch_df = df.iloc[start_idx:end_idx]
                
                # 创建指标对象
                indicator_objects = []
                
                for idx, row in batch_df.iterrows():
                    # 仅处理有效的每日行情ID
                    if pd.notna(row['daily_quote_id']):
                        # 创建技术指标实例
                        indicator = TechnicalIndicator(
                            daily_quote_id=int(row['daily_quote_id']),
                            ma5=row.get('ma5'),
                            ma10=row.get('ma10'),
                            ma20=row.get('ma20'),
                            ma30=row.get('ma30'),
                            ma60=row.get('ma60'),
                            ma120=row.get('ma120'),
                            ma250=row.get('ma250'),
                            vol5=row.get('vol5'),
                            vol10=row.get('vol10'),
                            rsi6=row.get('rsi6'),
                            rsi12=row.get('rsi12'),
                            rsi24=row.get('rsi24'),
                            macd=row.get('macd'),
                            macd_signal=row.get('macd_signal'),
                            macd_hist=row.get('macd_hist'),
                            boll_upper=row.get('boll_upper'),
                            boll_middle=row.get('boll_middle'),
                            boll_lower=row.get('boll_lower'),
                            kd_k=row.get('kd_k'),
                            kd_d=row.get('kd_d'),
                            kd_j=row.get('kd_j')
                        )
                        
                        indicator_objects.append(indicator)
                
                # 批量保存
                if indicator_objects:
                    self.db.add_all(indicator_objects)
                    self.db.commit()
                    saved_count += len(indicator_objects)
            
            return saved_count
            
        except Exception as e:
            self.db.rollback()
            logger.error(f"保存技术指标数据失败: {str(e)}")
            raise
    
    def calculate_all_securities_indicators(self, security_type=None, recalculate=False):
        """
        计算所有证券的技术指标
        
        Args:
            security_type: 证券类型，如'stock', 'index', 'etf'，默认为None表示所有类型
            recalculate: 是否重新计算所有指标，默认为False（仅计算新数据）
        
        Returns:
            计算的证券数量
        """
        # 查询证券列表
        query = self.db.query(SecurityInfo).filter(SecurityInfo.is_active == 1)
        
        if security_type:
            query = query.filter(SecurityInfo.type == security_type)
        
        securities = query.all()
        
        if not securities:
            logger.warning(f"没有找到需要计算指标的证券")
            return 0
        
        processed_count = 0
        error_count = 0
        
        # 遍历处理每个证券
        for security in securities:
            try:
                indicators_count = self.calculate_security_indicators(
                    security.code, recalculate
                )
                
                if indicators_count > 0:
                    processed_count += 1
                    logger.info(f"证券 {security.code} ({security.name}) 计算了 {indicators_count} 条指标")
                
            except Exception as e:
                error_count += 1
                logger.error(f"计算证券 {security.code} 的指标失败: {str(e)}")
        
        logger.info(f"完成指标计算，处理了 {processed_count} 个证券，失败 {error_count} 个")
        return processed_count


# 测试函数
def test_indicators_calculator():
    """测试技术指标计算器功能"""
    calculator = IndicatorsCalculator()
    
    # 从数据库中获取一个有足够数据的证券进行测试
    db = next(get_db())
    try:
        # 查找有足够日线数据的股票
        security = db.query(SecurityInfo).join(
            DailyQuote, SecurityInfo.id == DailyQuote.security_id
        ).group_by(SecurityInfo.id).having(
            func.count(DailyQuote.id) >= 30
        ).first()
        
        if security:
            print(f"使用证券 {security.code} ({security.name}) 进行测试")
            
            # 测试计算单个证券的技术指标
            indicators_count = calculator.calculate_security_indicators(security.code)
            print(f"成功计算 {indicators_count} 条技术指标")
            
        else:
            print("没有找到符合条件的证券数据")
    
    finally:
        db.close()


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 执行测试
    test_indicators_calculator()